# this is just to explore how to use the scipy.integrate.odeint function

from numpy import *
from pylab import *

# the bias current
Ie = lambda t: sin(2*pi+4*t*0.01)*cos(3*pi+4*t*0.01)

# transient sodium channel
vhalf_na = -35.0
na_inf = lambda v: 1/(1+exp((v-vhalf_na)/-k))
g_na = 20.0 # in mS/cm^2
E_na = 45.0 # in mV

# persistent potassium channel
vhalf_k = -36.0
k_inf = lambda v: 1/(1+exp((v-vhalf_k)/-k))
g_k = 4.5 # in mS/cm^2
E_k = -90.0 # in mV
tau_k = 0.6

# leak conductance
g_leak = 0.051 # in mS/cm^2
E_leak = -77.0 # in mV
k = 5

# cell membrane area and capacitance
C = 1.5 # uF/cm^2
# we'll ignore the area of the membrane
A = 1

def dX_dt(X,t=0):
    # X[0] is the membrane potential
    # X[1] is the potassium activation variable
    ret_array = empty(2,dtype=float32)
    ret_array[0] =  Ie(t) -g_na*na_inf(X[0])*(1-X[1])*(X[0]-E_na) -g_k*X[1]*(X[0]-E_k) -g_leak*(X[0]-E_leak) 
    ret_array[1] =(k_inf(X[0])-X[1])/tau_k
    return ret_array

def DX_dt(X,t=0):
    # X[0] is the membrane potential
    # X[1] is the potassium activation variable
    return array([ Ie(t) -g_na*(1/(1 + exp((X[0]-vhalf_na)/-k)))*(1-X[1])*(X[0]-E_na) -g_k*X[1]*(X[0]-E_k) -\
            g_leak*(X[0]-E_leak),((1/(1 + exp((X[0]-vhalf_k)/-k)))-X[1])/tau_k ]) 
    return ret_array


def odeint(func,init,time_vect):
    ''' integrate using the Runge-Kutta method:
        --func is a function who's first arguement is a list (or array) 
        of the dynamical variables (x) dx/dt=blah and second arguement is time.
        --init is a list (or array) of the initial values of the dynamical variables
        --time_vect is a list of time values.'''
    
    
    # first lets make everything numpy arrays
    results = empty( (len(time_vect),len(init)),dtype=float32 ) # reserve memory for an array of size t for our results
    time_vect = array( time_vect,dtype=float32 ) # just in case time isn't a numpy array, cast it as one
    init = array( init,dtype=float32 ) # same thing for the initialization array

    # now solve for the time course of the differential equation
    results[0] = init
    for i in range(1,len(time_vect)):
        delta = time_vect[i] - time_vect[i-1]
        t = time_vect[i-1]
        # calculate k1-k4
        k1 = func(results[i-1],t)
        k2 = func(results[i-1]+delta/2*k1,t+delta/2)    
        k3 = func(results[i-1]+delta/2*k2,t+delta/2)    
        k4 = func(results[i-1]+delta*k3,t+delta)  

        # calculate the new value of the dynamical variables
        results[i] = results[i-1] + delta/6 * (k1 + 2*k2 + 2*k3 + k4)  
    return results

# set the initial conditions for the integrator
#X0 = [-45,k_inf(-45)]
X0 = [-45,k_inf(-45)]

t = linspace(0,1000,100000) # take a million timesteps between 0 and 1 second

from scipy import integrate
#print "starting first integration"
#X = integrate.odeint(dX_dt,X0,t)
#print "doing first integration again"
X = integrate.odeint(dX_dt,X0,t)
#print "starting alternative integration"
#X= odeint(DX_dt,X0,t)
#print "finished integration"

from cell_test_functions import *
spikes = spike_detect(X[:,0], 0)

#X[range(20000, 60000), 0] = -60

plot(range(len(X[:,0])),X[:,0],spikes,X[spikes,0], 'r+') #,t,X[:,1])



